#ifndef __LINALG_MATRIX_H__
#define __LINALG_MATRIX_H__

#include "config.h"
#include "vector.h"


LINALG_BEGIN

template<typename T>

class Matrix : public Vector<T>
{
public:
    // Type definitions
    typedef T  value_type;
    typedef T* pointer;
    typedef T& reference;
    typedef const T* const_pointer;
    typedef const T& const_reference;
    typedef size_t size_type;

    typedef Vector<T> base_type;

public:
    // Constructors
    Matrix() : base_type(), nrow(0), ncol(0) {}

    Matrix(size_type _nrow, size_type _ncol) : base_type(_nrow * _ncol), nrow(_nrow), ncol(_ncol) { }

    Matrix(size_type _nrow, size_type _ncol, const_reference value) : base_type(_nrow * _ncol, value), nrow(_nrow), ncol(_ncol)  { }

    Matrix(const Matrix& other) : base_type(other), nrow(other.nrow), ncol(other.ncol) { }

    Matrix(Matrix&& other) noexcept : base_type(other), nrow(other.nrow), ncol(other.ncol) {}


public:
    Matrix& operator= (const Matrix& other)
    {
        if (this != &other)
        {
            base_type::operator=(other);
            nrow = other.nrow;
            ncol = other.ncol;
        }

        return *this;
    }


public:
    // Method
    void resize(size_type nrow, size_type ncol) { this->nrow = nrow; this->ncol = ncol; this->resize(nrow * ncol); }

    size_type rows() const { return nrow; }
    size_type cols() const { return ncol; }
    size_type size() const { return nrow * ncol; }

    void fill(const_reference value) { base_type::fill(value); }
    void fill_row(size_type i, const_reference value)
    {
        for(size_type j=0; j<ncol; ++j)
        {
            (*this)(i, j) = value;
        }
    }

    void fill_col(size_type j, const_reference value)
    {
        for(size_type i=0; i<nrow; ++i)
        {
            (*this)(i, j) = value;
        }
    }

    Matrix& from_value(const_reference _val)
    {
        resize(1, 1);
        (*this)[0] = _val;
        return *this;
    }

    void swap(Matrix& other) { std::swap(nrow, other.nrow); std::swap(ncol, other.ncol); base_type::swap(other); }

    void apply1(std::function<void(reference)>& f) { base_type::apply1(f); }

    void copy_from(const Matrix& other)
    {
        if (nrow != other.nrow || ncol != other.ncol)
        {
            throw std::invalid_argument("Matrix dimension mismatch");
        }

        base_type::copy_from(other);
    }

    void copy_from(pointer data, size_type size)
    {
        if (size!= nrow * ncol) { throw std::invalid_argument("Matrix dimension mismatch"); }
        
        base_type::copy_from(data, size);
    }

    reference operator() (size_type row, size_type col) { return base_type::operator[](row * ncol + col); }
    const_reference operator() (size_type row, size_type col) const { return base_type::operator[](row * ncol + col); }
    operator value_type () const { return (*this)[0]; }

    const Vector<T>& as_vector() const { return *this; }

public:
    friend std::ostream& operator<< (std::ostream& os, const Matrix& m)
    {
        os << "Matrix(" << m.nrow << ", " << m.ncol << ")" << std::endl;
        for (size_type i = 0; i < m.nrow; i++)
        {
            for (size_type j = 0; j < m.ncol; j++)
            {
                os << m(i, j) << " ";
            }
            os << std::endl;
        }
        return os;
    }


private:
    size_type nrow;
    size_type ncol;

    void allocate(size_type nrow, size_type ncol)
    {
        this->allocate(nrow * ncol);
    }
};

template<typename T>
Vector<T> operator* (const Matrix<T>& m, const Vector<T>& v)
{
    typedef typename Vector<T>::size_type size_type;

    if (m.cols()!= v.size()) { throw std::invalid_argument("Matrix dimension mismatch"); }

    size_type nrow = m.rows();
    size_type ncol = m.cols();

    Vector<T> result(nrow);

    for(size_type i=0; i<nrow; ++i)
    {
        T& r = result[i];
        for(size_type j=0; j<ncol; ++j)
        {
            r = Numeric<T>::add(result[i], Numeric<T>::mul(m(i, j), v[j]));
        }
    }

    return std::move(result);
}

template<typename T>
Vector<T> operator* (const Vector<T>& v, const Matrix<T>& m)
{
    typedef typename Vector<T>::size_type size_type;

    if (m.rows()!= v.size()) { throw std::invalid_argument("Matrix dimension mismatch"); }
    size_type nrow = m.rows();
    size_type ncol = m.cols();

    Vector<T> result(ncol);
    for(size_type i=0; i<ncol; ++i)
    {
        T& r = result[i];
        for(size_type j=0; j<nrow; ++j)
        {
            r = Numeric<T>::add(result[i], Numeric<T>::mul(m(i, j), v[j]));
        }
    }

    return std::move(result);
}


template<typename T>
class Numeric<Matrix<T> > : public Numeric<T>
{
public:
    static Matrix<T> add(const Matrix<T>& m1, const Matrix<T>& m2)
    {
        if (m1.rows()!= m2.rows() || m1.cols()!= m2.cols()) { throw std::invalid_argument("Matrix dimension mismatch"); }

        Matrix<T> m3(m1.rows(), m1.cols());
        for (size_t i = 0; i < m1.size(); i++)
        {
            m3[i] = Numeric<T>::add(m1[i], m2[i]);
        }

        return std::move(m3);
    }
    static Matrix<T> sub(const Matrix<T>& m1, const Matrix<T>& m2)
    {
        if (m1.rows()!= m2.rows() || m1.cols()!= m2.cols()) { throw std::invalid_argument("Matrix dimension mismatch"); }

        Matrix<T> m3(m1.rows(), m1.cols());
        for (size_t i = 0; i < m1.size(); i++)
        {
            m3[i] = Numeric<T>::sub(m1[i], m2[i]);
        }

        return std::move(m3);
    }

    // Matrix Multiply
    static Matrix<T> mul(const Matrix<T>& m1, const Matrix<T>& m2)
    {
        if (m1.cols() != m2.rows()) { throw std::invalid_argument("Matrix dimension mismatch"); }

        Matrix<T> m3(m1.rows(), m2.cols());

        for (size_t i = 0; i < m1.rows(); i++)
        {
            for (size_t j = 0; j < m2.cols(); j++)
            {
                for (size_t k = 0; k < m1.cols(); k++)
                {
                    m3(i, j) += m1(i, k) * m2(k, j);
                }
            }
        }

        return std::move(m3);
    }

};


namespace detail
{
    template<typename T>
    // 在矩阵A中找到第i列的最大元素，并将其和PIV[i]交换
    void find_max(Matrix<T>& A, Vector<typename Matrix<T>::size_type>& IPIV, typename Matrix<T>::size_type i)
    {
        // 定义矩阵A的size_type类型
        typedef typename Matrix<T>::size_type size_type;

        // 初始化最大元素的索引为i
        size_type max_index = i;
        // 遍历第i列之后的元素
        for(size_type j=i+1; j<A.rows(); ++j)
        {
            // 如果当前元素的绝对值大于最大元素的绝对值，则更新最大元素的索引
            if (Numeric<T>::abs(A(IPIV[j], i)) > Numeric<T>::abs(A(IPIV[max_index], i)))
            {
                max_index = j;
            }
        }

        // 交换最大元素和第i列的元素
        std::swap(IPIV[i], IPIV[max_index]);
    }

    // 以矩阵A第IPIV[i]行为基准，消去第IPIV[i+..]行
    template<typename T>
    void eliminate(
          Matrix<T>& A   // 左矩阵
        , Matrix<T>& B   // 右矩阵
        , typename Matrix<T>::size_type i  // 基准行索引
        , const Vector<typename Matrix<T>::size_type>& IPIV // 行索引映射表
        , Vector<T>& c)
    {
        // 定义矩阵A的size_type类型
        typedef typename Matrix<T>::size_type size_type;

        // 消元
        for(size_type j=i+1; j<A.rows(); ++j)
        {
            c[j] = Numeric<T>::div(A(IPIV[j], i), A(IPIV[i], i));
        }

        for(size_type j=i+1; j<A.rows(); ++j)
        {
            for(size_type k=i+1; k<A.cols(); ++k)
            {
                A(IPIV[j], k) -= Numeric<T>::mul(c[j], A(IPIV[i], k));
            }
            for(size_type k=0; k<B.cols(); ++k)
            {
                B(IPIV[j], k) -= Numeric<T>::mul(c[j], B(IPIV[i], k));
            }
        }
    }
}


// 基于选主元的高斯消元法（原位操作）
template<typename T>
void __mldivide(Matrix<T>& A, Matrix<T>& B)
{
    using namespace detail;

    typedef typename Matrix<T>::size_type size_type;

    if (A.cols() != B.rows()) { throw std::invalid_argument("Matrix dimension mismatch"); }

    size_type nrow = A.rows();
    size_type ncol = A.cols();
    size_type n = nrow;

    Vector<size_type> IPIV(n);
    for(size_type i=0; i<n; ++i)
    {
        IPIV[i] = i;
    }

    // 缓存消元系数
    Vector<T> c(n-1);

    // 消元
    for(size_type i=0; i<n-1; ++i)
    {
        // 查找主元
        find_max(A, IPIV, i);

        // 消元
        for(size_type j=i+1; j<n; ++j)
        {
            // 缓存消元基准行索引
            size_type s = IPIV[i];
            // 计算后续每一行的消元系数
            for(size_type k=i+1; k<n; ++k)
            {
                c[k] = Numeric<T>::div(A(IPIV[j], i), A(s, i));
            }

            // 消元
            eliminate(A, B, i, IPIV, c);
        }
    }

    // 回代
    for(size_type i=n-1; i>0; --i)
    {
        size_type s = IPIV[i];
        for(size_type j=0; j<i; ++j)
        {
            for(size_type k=0; k<B.cols(); ++k)
            {
                B(IPIV[j], k) -= Numeric<T>::mul(B(s, k), B(IPIV[i], k));
            }
        }
    }
}

template<typename T>
Matrix<T> mldivide(Matrix<T> A, Matrix<T> B)
{
    __mldivide(A, B);
    return std::move(B);
}

template<typename T>
Matrix<T> operator+ (const Matrix<T>& m1, const Matrix<T>& m2) { return Numeric<Matrix<T> >::add(m1, m2); }

template<typename T>
Matrix<T> operator- (const Matrix<T>& m1, const Matrix<T>& m2) { return Numeric<Matrix<T> >::sub(m1, m2); }

template<typename T>
Matrix<T> operator* (const Matrix<T>& m1, const Matrix<T>& m2) { return Numeric<Matrix<T> >::mul(m1, m2); }

LINALG_END
#endif